from neo4j import GraphDatabase
import ast
import warnings
from openai import OpenAI

class TupleToNeo4j:
    """
    这是一个将元组提交到Neo4j数据库的类
    """
    # def __init__(self,username:str,password:str,uri:str="bolt://localhost:7687"):
    #     """
    #     初始化TupleToNeo4j类
    #     :param username: Neo4j用户名
    #     :param password: Neo4j密码
    #     :param uri: Neo4j连接地址，默认为：'bolt://localhost:7687'
    #     """
    #     self.driver = GraphDatabase.driver(uri, auth=(username, password))
    #
    #
    def __init__(self,driver: GraphDatabase.driver):
        self.driver = driver

    def __createNode(self, node_tuple, isStartNode=None):
        """
        创建节点的CQL语句
        :param node_tuple: 节点元组，格式为：(节点标签,{属性名:属性值})
        :param isStartNode: 是否是起始节点，默认为None。None:单独节点，Ture:开始节点,False:结束节点
        :return: 节点的CQL语句
        """

        # 判断当前节点是开始还是结束节点
        # 如果是起始节点，则创建起始节点的CQL语句。
        # 如果是结束节点，则创建结束节点的CQL语句。
        if isStartNode is None:
            cql = "merge(n:%(lableName)s{" % {"lableName": node_tuple[0]}
        elif isStartNode:
            cql = "merge(start:%(lableName)s{" % {"lableName": node_tuple[0]}
        else:
            cql = "merge(end:%(lableName)s{" % {"lableName": node_tuple[0]}

        # 获取节点所有属性名和属性值
        keylist = list(node_tuple[1].keys())
        valuelist = list(node_tuple[1].values())

        # 初始化节点的属性名和属性值，满足CQL语法
        for i in range(len(keylist)):
            # 如果属性值是字符串，则加上引号。
            if type(valuelist[i]) == str:
                valuelist[i] = "'" + valuelist[i] + "'"

            # 如果属性值不是字符串，则转换为字符串形式后续进行拼接。
            if type(valuelist[i]) != str:
                valuelist[i] = str(valuelist[i])

            # 如果属性名不是字符串，则转换为字符串形式后续进行拼接。
            if type(keylist[i]) != str:
                valuelist[i] = str(valuelist[i])

            # 将属性名和属性值与创建节点的CQL语句进行拼接。
            # 如果是最后一个属性，则不加逗号，反之加逗号。
            if i == len(keylist) - 1:
                cql += keylist[i] + ":" + valuelist[i]
            else:
                cql += keylist[i] + ":" + valuelist[i] + ","
        cql += "})\n"

        return cql

    def __createRelationAndNodeCQl(self, dataTuple, showCqlCode=False):
        """
        创建节点和关系CQL语句
        :param dataTuple: 三元组,格式为：((节点标签,{属性名:属性值}),关系标签,(节点标签,{属性名:属性值}))
        :param showCqlCode: 是否展示拼接完成后的CQL命令
        :return: CQL语法
        """
        # 获取起始节点和结束节点的元组
        startNodeTuple = dataTuple[0]
        startNodeCQL = self.__createNode(startNodeTuple, True)

        endNodeTuple = dataTuple[2]
        endNodeCQL = self.__createNode(endNodeTuple, False)

        # 获取关系标签和关系CQL语句
        relationLable = dataTuple[1]
        relationCQL = "create (start)-[r:%(relationLable)s]->(end)" % {"relationLable": relationLable}

        # 拼接节点和关系CQL语句
        createRelationAndNodeCQL = startNodeCQL + endNodeCQL + relationCQL

        # 展示拼接后的CQL命令
        if showCqlCode:
            print("=" * 25)
            print(createRelationAndNodeCQL)
            print("=" * 25)
        return createRelationAndNodeCQL

    def __deleteAllNodeAndReletion(self, isVerify=True, showLog=False):
        """
        清空Neo4j中所有节点和关系
        :param isVerify: 是否删除前确认
        :param showLog: 是否展示删除后的数据库
        :return: 无返回值
        """
        if isVerify:
            print("即将清空数据库是否清空")
            if input("是则请输入'Y'\n[User]:") == "Y":
                self.driver.execute_query("match(n) detach delete n")
                if showLog:
                    print("数据库已清空\n")
                    print("当前数据库所有内容：")
                    print(self.driver.execute_query("match(n) return n"))
        else:
            self.driver.execute_query("match(n) detach delete n")
            if showLog:
                print("数据库已清空\n")
                print("当前数据库所有内容：")
                print(self.driver.execute_query("match(n) return n"))

    def createRelationAndNode(self, dataTuple, showCqlCode=False):
        """
        执行创建节点和关系
        :param dataTuple: 三元组,格式为：((节点标签,{属性名:属性值}),关系标签,(节点标签,{属性名:属性值}))
        :param showCqlCode: 是否展示拼接完成后的CQL命令
        :return: 无返回值
        """
        self.driver.execute_query(self.__createRelationAndNodeCQl(dataTuple, showCqlCode))

    def createRelationsAndNodes(self, dataTuples, cover=True, showLog=False):
        """
        批量创建节点和关系
        :param dataTuples: 三元组列表,格式为：[((节点标签,{属性名:属性值}),关系标签,(节点标签,{属性名:属性值})),((节点标签,{属性名:属性值}),关系标签,(节点标签,{属性名:属性值})),...]
        :param cover: 是否覆盖数据库中的数据,(默认)Ture:覆盖数据库，False：不覆盖
        :param showLog: 是否展示数据库中的数据,Ture:展示,(默认):不展示
        :return: 无返回值
        """
        if cover:
            self.__deleteAllNodeAndReletion(False)
        for dataTuple in dataTuples:
            self.createRelationAndNode(dataTuple=dataTuple, showCqlCode=showLog)
        print("创建关系成功")


class DataTuple:
    """
    这是一个将元组格式的字符串转成元组的类
    """
    def getTupleListFromFile(splitCode, filePath):
        """
        将某一个文件内容转成元组
        :param splitCode: 分割符
        :param filePath: 文件地址
        :return: 元组列表
        """
        lineData = open(filePath, 'r', encoding='utf-8').read()

        strTuples = lineData.split(splitCode)

        tupleList = []

        for strTuple in strTuples:
            tupleList.append(ast.literal_eval(strTuple))

        return tupleList

    def getTupleListFromVar(splitCode, textSource):
        """
        将字符串转成元组
        :param splitCode: 分割符
        :param textSource: 数据源
        :return: 元组列表
        """
        lineData = textSource

        strTuples = lineData.split(splitCode)

        tupleList = []

        for strTuple in strTuples:
                tupleList.append(ast.literal_eval(strTuple))
        return tupleList

class Presenter:
    """
    发言者枚举类
    """
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"

class ConstanstLLM():
    """
    负责存放LLM常量的类
    """
    # 提示词---提示词优化器
    prompt_word_optimizer = """
    【提示词优化器】
    "作为专业提示词工程师，请根据以下原则优化我的原始提示词：
    明确目标：清晰定义任务目的与预期成果结构化表达：使用分点/编号等逻辑结构关键词强化：突出核心概念与限定条件去除冗余：删除重复/模糊表述场景适配：补充必要上下文信息格式规范：统一术语与表述方式
    请仅返回优化后的最终版本，无需解释过程。示例输入：'写一篇关于环保的文章' → 输出：'以专业视角撰写2000字深度文章，系统分析当前环境治理政策与技术创新，需包含案例实证与数据支撑，采用学术论文格式。'
    现在请优化以下提示词：[你的原始提示词]"
    （使用时将方括号内容替换为你的原始提示词即可）
    """

    # 提示词---内容总结
    content_summarizer = """
    [文本内容]
    请按照以下要求总结：
    1. 使用"是一个关于xx的文本"的固定句式
    2. xx部分应：
        - 准确反映文本核心主题
        - 控制在5-10个汉字内
        - 保持客观中立
    """

    # 提示词---需求描述
    def getPromptWord(lineWord):
        prompt_word = """
        下述文本%s,我需要将文本所有内容提取节点和节点之间的关系，
        通过创建一个存储开始结束节点元组和关系的三元组来存储图数据，
        元组的元素存放开始和结束节点的标签与属性和两个元组之间关系标签 示例结构：((节点标签,{属性名:属性值}),"关系标签",(节点标签,{属性名:属性值})),
        每个三元组之间使用英文的分号(;)隔开不换行,如果可以每个节点属性尽量控制在一个，同时属性名固定为"name"
        ，且一段文本尽可能多的提取元组，标签和关系都使用中文，最后要确保图中任意一个节点一定能到达另一个任意的节点,
        同时要满足第一范式,只返回三元组结果,解释性文本一概不返回，示例返回格式：(("Person",{"name":"张三"}),"认识",("Person",{"name":"李四"));(("Person",{"name":"张三"}),"年龄",("age",{name:23}))
        """ % lineWord
        return prompt_word

    # 返回Message
    def getPrompt(mode, prompt):
        """
        格式化提示词
        :param mode:发言者
        :param prompt:提示词
        :return: message字典
        """
        warnings.warn("该方法不推荐使用，请使用getMessage()方法获取提示词Message")
        return {"role": mode, "content": prompt}


    def getMessage(presenter:str, prompt):
        """
        格式化提示词
        :param presenter:发言者
        :param prompt: 提示词
        :return: message字典
        """
        return {"role": presenter, "content": prompt}


    # # 获取被提取的文本内容
    # def getFileText(filePath):
    #     # 获取指定文件文本内容
    #     return open(filePath, "r", encoding="utf-8").read()
    # def getStringText(sourceText):
    #     # 获取指定字符串文本内容
    #     return sourceText


class CreateStrTupleByLLM():
    """
    这是通过LLM提取一段文本内容中的关系，并将其输出为三元组格式的字符串的类
    """
    # 是否展示日志
    # showLog=False

    def __init__(self,client,showLog):
        """
        初始化client
        :param client:
        """
        self.client = client
        self.showLog=showLog

    def __createClient(api_key,  base_url):
        """
        创建OpenAI对象
        :param api_key: LLM的api_key
        :param base_url: LLM的base_url
        :return: OpenAI对象
        """
        warnings.warn("该方法为开发时不规范写法，生成环境请自行创建。")
        return OpenAI(api_key=api_key,base_url=base_url)

    def __getPromptWordOptimizer(self,prompt:str):
        """
        负责优化提示词
        :param prompt:被提取的文本
        :return: 优化后的提示词
        """
        client=self.client
        messages=[]
        # 将助手添加到消息列表中
        messages.append(ConstanstLLM.getMessage(Presenter.SYSTEM,ConstanstLLM.prompt_word_optimizer))

        # 将需要被优化的提示词加入到消息列表中
        system_prompt="请帮我优化提示词：'"+ConstanstLLM.getPromptWord(lineWord=self.__getContentSummarizer(prompt))+"'"
        messages.append(ConstanstLLM.getMessage(Presenter.USER,system_prompt))

        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
        )

        # 打印优化后的提示词
        if  self.showLog:
            print(response.choices[0].message.content)
        return response.choices[0].message.content

    def __getContentSummarizer(self,prompt:str):
        """
        获取内容总结
        :param prompt:被总结的提示词
        :return:总结后的提示词
        """
        client=self.client
        messages=[]
        # 将助手添加到消息列表中
        messages.append(ConstanstLLM.getMessage(Presenter.SYSTEM,ConstanstLLM.content_summarizer))
        # 将需要被总结的提示词加入到消息列表中
        messages.append(ConstanstLLM.getMessage(Presenter.USER,prompt))

        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
        )
        # 打印总结后的提示词
        if  self.showLog:
            print(response.choices[0].message.content)
        return response.choices[0].message.content

    def getTuple(self,prompt:str):
        """
        获取三元组
        :param prompt:被提取的文本内容
        :return:三元组
        """
        client=self.client
        messages=[]
        # 将助手添加到消息列表中
        messages.append(ConstanstLLM.getMessage(Presenter.SYSTEM,self.__getPromptWordOptimizer(prompt)))
        # 将需要被提取的文本内容加入到消息列表中
        messages.append(ConstanstLLM.getMessage(Presenter.USER,prompt))

        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
        )
        #  打印三元组
        if  self.showLog:
            print(response.choices[0].message.content)

        return response.choices[0].message.content